from typing import TypedDict, List
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END
from IPython.display import Image, display
import sqlite3

import langgraph
print(dir(langgraph))

from langgraph.checkpoint.sqlite import SqliteSaver



from langchain_openai import ChatOpenAI

import os
# ================== 定义上下文-状态结构 ==================
class AgentState(TypedDict):
    user_id: str
    message: str
    intent: str  #意图 register / faq / off_topic
    current_step: str
    is_registered: bool
    collected_info: dict
    context: List[str]  # 对话历史

# ================== 初始化 LLM ==================
def create_llm() -> ChatOpenAI:
    """创建并返回配置好的大模型实例"""
    return ChatOpenAI(
        api_key=os.getenv("DASHSCOPE_API_KEY") or "your_api_key_here",
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        model="qwen2.5-math-72b-instruct"
    )

llm = create_llm()

# ================== 节点函数定义 ==================

def greet_user(state: AgentState) -> AgentState:
    return {
        **state,
        "message": "你好！我是你的专属助手，有什么可以帮您的吗？",
        "current_step": "greet"
    }

def detect_intent(state: AgentState) -> AgentState:
    prompt = f"""
你是一个智能客服助手，请分析以下用户消息的意图：
"{state['message']}"

请判断：
- 是否与注册流程有关？
- 是否是常见问题咨询（如价格、联系方式）？
- 是否完全偏离流程？

返回 intent 类型（register/faq/off_topic）
"""
    response = llm.invoke(prompt)
    intent = response.content.strip().lower()
    return {**state, "intent": intent}

def guide_to_register(state: AgentState) -> AgentState:
    if not state["is_registered"]:
        return {
            **state,
            "message": "您还未注册哦，是否需要我帮您快速完成注册流程？",
            "current_step": "guide_to_register"
        }
    else:
        return {**state, "message": "欢迎回来！"}

def collect_user_info(state: AgentState) -> AgentState:
    info = {"email": "user@example.com"}  # 示例数据，实际可由对话提取
    return {
        **state,
        "collected_info": info,
        "is_registered": True,
        "message": "感谢注册！我们将尽快联系您。",
        "current_step": "collect_user_info"
    }

def offer_next_steps(state: AgentState) -> AgentState:
    return {
        **state,
        "message": "这是您的优惠券码：WELCOME2025。有任何问题欢迎随时联系我！",
        "current_step": "offer_next_steps"
    }

def handle_off_topic(state: AgentState) -> AgentState:
    prompt = f"""
用户说：“{state['message']}”

请根据你的知识库回答该问题。如果你不确定答案，请说明“我暂时无法提供准确信息”。

注意：请保持简洁明了。
"""
    response = llm.invoke(prompt)
    reply = response.content.strip()
    return {
        **state,
        "message": reply,
        "intent": "faq" if "contact" in reply or "price" in reply else "off_topic"
    }

def resume_flow(state: AgentState) -> AgentState:
    return {
        **state,
        "message": f"刚刚我们说到哪里了？让我们继续吧！当前步骤：{state['current_step']}"
    }

# ================== 条件分支逻辑 ==================
def route_after_greet(state: AgentState):
    if "register" in state["intent"].lower():
        return "guide_to_register"
    elif state["intent"] in ["faq", "off_topic"]:
        return "handle_off_topic"
    else:
        return "offer_next_steps"

def route_after_handle_off_topic(state: AgentState):
    return state.get("current_step", "greet")  # 返回上一个流程节点

# ================== 构建状态图 ==================
workflow = StateGraph(AgentState)

# 添加节点
workflow.add_node("greet_user", greet_user)
workflow.add_node("detect_intent", detect_intent)
workflow.add_node("guide_to_register", guide_to_register)
workflow.add_node("collect_user_info", collect_user_info)
workflow.add_node("offer_next_steps", offer_next_steps)
workflow.add_node("handle_off_topic", handle_off_topic)
workflow.add_node("resume_flow", resume_flow)

# 设置入口点
workflow.set_entry_point("greet_user")

# 构建流程
workflow.add_edge("greet_user", "detect_intent")
workflow.add_conditional_edges(
    "detect_intent",
    route_after_greet,
    {
        "guide_to_register": "guide_to_register",
        "handle_off_topic": "handle_off_topic",
        "offer_next_steps": "offer_next_steps"
    }
)
workflow.add_edge("guide_to_register", "collect_user_info")
workflow.add_edge("collect_user_info", "offer_next_steps")
workflow.add_edge("offer_next_steps", END)

# 非流程问题处理路径
workflow.add_conditional_edges(
    "handle_off_topic",
    route_after_handle_off_topic,
    {
        "greet": "greet_user",
        "guide_to_register": "guide_to_register",
        "collect_user_info": "collect_user_info",
        "offer_next_steps": "offer_next_steps"
    }
)

# 编译图谱 + 启用 checkpointer（用于记忆会话）
# memory = SqliteSaver.from_conn_string(":memory:")  # 可改为文件路径或数据库

sqlite_conn = sqlite3.connect("checkpoints.sqlite", check_same_thread=False)
checkpointer = SqliteSaver(sqlite_conn)
app = workflow.compile(checkpointer=checkpointer)

print(f"[DEBUG] Checkpointer 类型：{type(checkpointer)}")
# Draw the graph
try:
    app.get_graph(xray=True).draw_mermaid_png(output_file_path="../graph.png")
except Exception:
    pass
# ================== 多轮对话主循环 ==================
if __name__ == "__main__":
    import uuid

    print("启动客户拉新机器人（输入 'exit' 退出）...\n")

    # 每个用户使用一个 thread_id 来保存状态
    user_id = input("请输入用户ID（例如：u123456）: ")
    config = {"configurable": {"thread_id": user_id}}

    try:
        # 获取初始状态（如果存在），赋值给上下文
        initial_state = app.get_state(config).values
    except:
        initial_state = {
            "user_id": user_id,
            "message": "",
            "intent": "",
            "current_step": "",
            "is_registered": False,
            "collected_info": {},
            "context": []
        }

    while True:
        user_input = input("\n[用户]: ").strip()
        if user_input.lower() in ["exit", "quit", "q"]:
            print("再见！")
            break

        # 更新 message 字段
        initial_state["message"] = user_input

        # 运行一次图谱
        result = app.invoke(initial_state, config)

        # 打印机器人回复
        print(f"[机器人]: {result['message']}")

        # 更新上下文和状态
        initial_state = result